import numpy as np
import heapq

# 1. 数据生成
def generate_data(num_items=50, dimensions=5):
    """ 
    生成内容特征向量和用户特征向量
    :param num_items: 内容数量
    :param dimensions: 向量维度
    :return: 内容特征矩阵，用户特征向量
    """ 
    np.random.seed(42)
    item_vectors = np.random.rand(num_items, dimensions)
    user_vector = np.random.rand(1, dimensions) # 单个用户的兴趣特征
    return item_vectors, user_vector

item_vectors, user_vector = generate_data(num_items=50, dimensions=5)

# 2. HNSW索引实现
class HNSW:
    """ 
    实现HNSW用于近邻搜索
    """ 
    def __init__(self, vectors, max_neighbors=5, max_layers=3):
        self.vectors = vectors
        self.max_neighbors = max_neighbors
        self.max_layers = max_layers
        self.layers = self._construct_layers()
    
    def _construct_layers(self):
        layers = []
        current_vectors = self.vectors
        for layer_id in range(self.max_layers):
            layer = {i: [] for i in range(len(current_vectors))}
            for i in range(len(current_vectors)):
                distances = []
                for j in range(len(current_vectors)):
                    if i != j:
                        dist = np.linalg.norm(current_vectors[i] - current_vectors[j])
                        distances.append((dist, j))
                
                nearest_neighbors = heapq.nsmallest(self.max_neighbors, distances)
                layer[i] = [neighbor[1] for neighbor in nearest_neighbors]
            layers.append(layer)
            current_vectors = current_vectors[::2] # 每层点数减半
        return layers

    def search(self, query_vector):
        """ 
        在HNSW中搜索与查询向量最相似的内容
        :param query_vector: 查询向量
        :return: 最相似的内容索引和距离
        """
        current_node = 0
        for layer in range(len(self.layers) - 1, -1, -1):
            layer_graph = self.layers[layer]
            layer_vectors = self.vectors[:len(layer_graph)]
            min_distance = np.linalg.norm(query_vector - layer_vectors[current_node])
            nearest_node = current_node
            while True:
                updated = False
                for neighbor in layer_graph[current_node]:
                    distance = np.linalg.norm(query_vector - layer_vectors[neighbor])
                    if distance < min_distance:
                        min_distance = distance
                        nearest_node = neighbor
                        updated = True
                if not updated:
                    break
                current_node = nearest_node
        return nearest_node, min_distance

# 3. 构建HNSW索引
hnsw_index = HNSW(item_vectors, max_neighbors=5, max_layers=3)

# 4. 查询推荐内容
def recommend_items(user_vector, hnsw_index, top_k=5):
    """ 
    为用户推荐与其兴趣最相似的内容
    :param user_vector: 用户兴趣特征向量
    :param hnsw_index: HNSW索引
    :param top_k: 推荐内容数量
    :return: 推荐内容索引和相似度
    """
    recommendations = []
    for _ in range(top_k):
        nearest_node, distance = hnsw_index.search(user_vector)
        recommendations.append((nearest_node, distance))
    return recommendations

# 5. 输出结果
recommendations = recommend_items(user_vector, hnsw_index, top_k=5)

print("用户特征向量:")
print(user_vector)
print("\n内容特征向量:")
print(item_vectors)
print("\n推荐内容:")
for idx, (item_idx, distance) in enumerate(recommendations):
    print(f"推荐{idx + 1}: 内容索引{item_idx}, 相似度距离{distance:.6f}")